package map;

import java.util.concurrent.locks.ReentrantReadWriteLock;

public class MyConcurrentHashMap<K, V> {

    private static final int NUM_BUCKETS = 16; // 桶的数量
    private Node<K, V>[] buckets; // 存储键值对的数组
    private ReentrantReadWriteLock[] locks; // 存储锁的数组

    public MyConcurrentHashMap() {
        buckets = new Node[NUM_BUCKETS];
        locks = new ReentrantReadWriteLock[NUM_BUCKETS];
        for (int i = 0; i < NUM_BUCKETS; i++) {
            locks[i] = new ReentrantReadWriteLock();
        }
    }

    private int hash(K key) {
        return Math.abs(key.hashCode()) % NUM_BUCKETS;
    }

    public V put(K key, V value) {
        int bucketIndex = hash(key);
        locks[bucketIndex].writeLock().lock(); // 获得写锁
        try {
            Node<K, V> current = buckets[bucketIndex];
            while (current != null) {
                if (current.key.equals(key)) {
                    V oldValue = current.value;
                    current.value = value;
                    return oldValue;
                }
                current = current.next;
            }
            Node<K, V> newNode = new Node<>(key, value);
            newNode.next = buckets[bucketIndex];
            buckets[bucketIndex] = newNode;
            return null;
        } finally {
            locks[bucketIndex].writeLock().unlock(); // 释放写锁
        }
    }

    public V get(K key) {
        int bucketIndex = hash(key);
        locks[bucketIndex].readLock().lock(); // 获得读锁
        try {
            Node<K, V> current = buckets[bucketIndex];
            while (current != null) {
                if (current.key.equals(key)) {
                    return current.value;
                }
                current = current.next;
            }
            return null;
        } finally {
            locks[bucketIndex].readLock().unlock(); // 释放读锁
        }
    }

    public V remove(K key) {
        int bucketIndex = hash(key);
        locks[bucketIndex].writeLock().lock(); // 获得写锁
        try {
            Node<K, V> current = buckets[bucketIndex];
            Node<K, V> previous = null;
            while (current != null) {
                if (current.key.equals(key)) {
                    if (previous == null) {
                        buckets[bucketIndex] = current.next;
                    } else {
                        previous.next = current.next;
                    }
                    return current.value;
                }
                previous = current;
                current = current.next;
            }
            return null;
        } finally {
            locks[bucketIndex].writeLock().unlock(); // 释放写锁
        }
    }

    private static class Node<K, V> {
        K key;
        V value;
        Node<K, V> next;

        public Node(K key, V value) {
            this.key = key;
            this.value = value;
            this.next = null;
        }
    }
}

